import cv2
import os
import argparse

from mediapipe_utils import MediaPipeUtils
from detectron2_utils import Detectron2Utils


class PoseRecognition:
    def __init__(self):
        model_path = '../data/pose_landmarker_heavy.task'
        if not os.path.exists(model_path):
            print(f"Model file not found at {model_path}")
            return
        running_mode = 'VIDEO'
        num_poses = 10
        self.mediapipe_utils = MediaPipeUtils(model_path, running_mode, num_poses)
        self.detectron2_utils = Detectron2Utils()

    def process_image(self, image_path, output_filename, scale_percent=100):
        # 调用 detectron2_utils 的 body_posture_backbone_image 方法获取处理后的图片
        d2_utils = Detectron2Utils()
        processed_image_path = d2_utils.body_posture_backbone_image(image_path, output_filename, scale_percent)
        # 使用 OpenCV 显示处理后的图片
        processed_image = cv2.imread(processed_image_path)
        processed_image = resize_frame(processed_image, scale_percent)
        cv2.imshow('Processed Image', processed_image)
        cv2.waitKey(0)
        cv2.destroyAllWindows()

    def process_video(self, video_path, scale_percent=100):
        """
        处理视频文件，实现每帧的姿态识别。
        Args:
            video_path (str): 视频文件路径
            scale_percent(int): 定义缩小的比例，例如缩小到原图的60%
        """
        cap = cv2.VideoCapture(video_path)
        if not cap.isOpened():
            print(f"Failed to open video file at {video_path}")
            return

        prev_landmarks = None
        while cap.isOpened():
            ret, frame = cap.read()
            if not ret:
                print("End of video")
                break

            timestamp_ms = int(cap.get(cv2.CAP_PROP_POS_MSEC))

            # 使用 MediaPipe 进行姿态检测
            frame = resize_frame(frame, scale_percent)
            res = self.mediapipe_utils.detect_for_video(frame, timestamp_ms)

            if res.pose_landmarks:
                # 绘制所有姿态关键点
                annotated_image = self.mediapipe_utils.draw_landmarks_on_image(frame, res)
                # 更新前一帧的关键点
                prev_landmarks = res.pose_landmarks[0] if res.pose_landmarks else None
            else:
                annotated_image = frame  # 如果没有检测到姿态，直接返回原图

            # 显示处理后的图像
            cv2.imshow('Pose Estimation', annotated_image)

            if cv2.waitKey(1) & 0xFF == ord('q'):
                break

        cap.release()
        cv2.destroyAllWindows()


def resize_frame(frame, scale_percent=100):
    width = int(frame.shape[1] * scale_percent / 100)
    height = int(frame.shape[0] * scale_percent / 100)
    dim = (width, height)
    resized_frame = cv2.resize(frame, dim, interpolation=cv2.INTER_AREA)
    return resized_frame


def parse_arguments():
    parser = argparse.ArgumentParser(description="Pose Recognition Script")

    # 设置互斥组，确保只能选择视频或图片模式
    mode_group = parser.add_mutually_exclusive_group(required=True)
    mode_group.add_argument('--video', '-v', help='Path to the video file')
    mode_group.add_argument('--image', '-i', help='Path to the image file')

    # 公共参数
    parser.add_argument('--scale', '-s', type=int, default=100, help='Scale percentage')

    # 图片模式特有的参数
    parser.add_argument('--output', '-o', help='Output filename for images (including extension)')

    return parser.parse_args()


if __name__ == '__main__':
    os.environ['OMP_NUM_THREADS'] = '1'

    pose_recognition = PoseRecognition()

    """
    两种调用方式(scale可以不加)：
    1，python pose_recognition.py --video ../data/video/rec-8.mp4 --scale 30
    2，python pose_recognition.py --image ../data/image/nsr.jpg --output nsr_pr.jpg --scale 30
    """
    args = parse_arguments()
    if args.video:pose_recognition.process_video(args.video, args.scale)
    elif args.image:
        if not args.output:
            raise ValueError("Output filename is required when processing an image.")
        pose_recognition.process_image(args.image, args.output, args.scale)
